# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import datetime
import logging
import sys

from coin.base.datetime_util import parse_end_time
from coin.exchange.base.ws import writer
from coin.exchange.base.ws.subscriber import RotatingSubscriber
from coin.exchange.okex import config
from coin.exchange.okex_futures.ws.subscriber import OkexFuturesSubscriber


class BaseSubscriberFactory(object):
  def __init__(self, group_id, worker_id, logger=None, futures_group_symbol=None):
    self._group_id = group_id
    self._worker_id = worker_id
    self._logger = logger
    self._group_symbol = futures_group_symbol or config.futures_group_symbol

  def init_subscriber(self, queue, ioloop, close_callback):
    sub = OkexFuturesSubscriber(queue,
                                self._group_id,
                                self._worker_id,
                                self._group_symbol[self._group_id],
                                close_callback=close_callback,
                                logger=self._logger,
                                subscribe_ticker=True)
    sub.start(ioloop)
    return sub


def main(argv):
  assert len(argv) >= 2, '<worker_id> is missing.'
  worker_id = argv[1].strip()
  if len(argv) == 3:
    groups = sorted(set([c for c in argv[2]]))
  else:
    groups = sorted(config.futures_group_symbol.keys())
  logging.info('Subscribing groups: %s' % ', '.join(groups))

  unit_td = datetime.timedelta(hours=1)
  if int(worker_id) % 2 == 1:
    base_offset_td = datetime.timedelta(minutes=0)
    end_time = parse_end_time('23:59:00')
  else:
    base_offset_td = datetime.timedelta(minutes=30)
    end_time = parse_end_time('00:29:00')

  subscribers = []

  def init(queue, ioloop):
    for group_id in groups:
      group_idx = ord(group_id) - ord('a')
      offset_td = base_offset_td + datetime.timedelta(seconds=5 * group_idx)

      logger = logging.getLogger('okex_ws_futures_subscriber_%s%s' % (group_id, worker_id))
      factory = BaseSubscriberFactory(group_id, worker_id, logger)
      sub = RotatingSubscriber(factory.init_subscriber, queue, unit_td, offset_td, logger=logger)
      sub.start(ioloop, datetime.timedelta(seconds=5 * group_idx))
      subscribers.append(sub)

  def teardown():
    for s in subscribers:
      s.stop()

  writer.run2(init, teardown, end_time=end_time)


if __name__ == '__main__':
  logging.basicConfig(level='DEBUG', format='%(levelname)s %(asctime)s %(name)s] %(message)s')
  sys.exit(main(sys.argv))
